
import unittest
from src.core.Commander import Commander

from src.core.Proxy import CallerProxy, DataFrameProxy, PandasProxy
import pandas as pd


class DataFrameProxyTest(unittest.TestCase):

    def test_getItem(self):
        df = pd.DataFrame({
            'name': 'x1 x2 x3 x4'.split(),
            'n1': [1, 2, 3, 4],
            'n2': [1, 2, 3, 4],
        })

        cp = CallerProxy()
        cp.set_obj_to_var('df', df)

        with cp.with_cmd() as cmd:
            cmd.create_statement('df', 'df')
            cp['name']

        expected = "df = df['name']"
        act = cp.to_code()

        self.assertEqual(expected, act)

    def test_getItem_and_other_call(self):
        df = pd.DataFrame({
            'name': 'x1 x2 x3 x4'.split(),
            'n1': [1, 2, 3, 4],
            'n2': [1, 2, 3, 4],
        })

        cp = CallerProxy()
        cp.set_obj_to_var('df', df)

        with cp.with_cmd() as cmd:
            cmd.create_statement('df', 'df')
            cp['name'].isin([1, 2, 3])

        expected = "df = df['name'].isin([1, 2, 3])"
        act = cp.to_code()

        self.assertEqual(expected, act)

    def test_groupby_without_ret(self):
        df = pd.DataFrame({
            'name': 'x1 x2 x3 x4'.split(),
            'n1': [1, 2, 3, 4],
            'n2': [1, 2, 3, 4],
        })

        cp = CallerProxy()
        cp.set_obj_to_var('df', df)

        with cp.with_cmd() as cmd:
            cmd.create_statement('df', 'df')
            cp.groupby('name')
            cp.agg({'n1': ['sum']})

        expected = "df = df.groupby('name').agg({'n1': ['sum']})"
        act = cp.to_code()

        self.assertEqual(expected, act)

    def test_groupby_with_ret(self):
        df = pd.DataFrame({
            'name': 'x1 x2 x3 x4'.split(),
            'n1': [1, 2, 3, 4],
            'n2': [1, 2, 3, 4],
        })

        cp = CallerProxy()
        cp.set_obj_to_var('df', df)

        with cp.with_cmd() as cmd:
            cmd.create_statement('df', 'gp')
            cp.groupby('name')
            cp.agg({'n1': ['sum']})

        expected = "gp = df.groupby('name').agg({'n1': ['sum']})"
        act = cp.to_code()

        self.assertEqual(expected, act)

    def test_2_cmds(self):
        df = pd.DataFrame({
            'name': 'x1 x2 x3 x4'.split(),
            'n1': [1, 2, 3, 4],
            'n2': [1, 2, 3, 4],
        })

        cp = CallerProxy()
        cp.set_obj_to_var('df', df)

        with cp.with_cmd() as cmd:
            cmd.create_statement('df', 'df_where')
            cp.query('name=="x1"')

            cmd.create_statement('df_where', 'gp')
            cp.groupby('name')
            cp.agg({'n1': ['sum']})

        exp1 = '''df_where = df.query('name=="x1"')'''
        exp2 = "gp = df_where.groupby('name').agg({'n1': ['sum']})"
        expected = '\n'.join([exp1, exp2])

        act = cp.to_code()

        self.assertEqual(expected, act)

    def test_dataframe(self):
        df = pd.DataFrame({
            'name': 'x1 x2 x3 x4'.split(),
            'n1': [1, 2, 3, 4],
            'n2': [1, 2, 3, 4],
        })

        cp = CallerProxy()
        cp.set_obj_to_var('df', df)

        with cp.with_cmd() as cmd:
            cmd.create_statement('df', 'df_where')
            cp.query('name=="x1"')

            cmd.create_statement('df_where', 'gp')
            cp.groupby('name')
            cp.agg({'n1': ['sum']})

        df_ret = cp.run('gp')
        exp_df = df.query('name=="x1"').groupby('name').agg({'n1': ['sum']})

        self.assertTrue(exp_df.equals(df_ret))

    def test_dataframe_2_cmds(self):
        df = pd.DataFrame({
            'name': 'x1 x2 x3 x4'.split(),
            'n1': [1, 2, 3, 4],
            'n2': [1, 2, 3, 4],
        })

        cp = CallerProxy()
        cp.set_obj_to_var('df', df)

        with cp.with_cmd() as cmd:
            cmd.create_statement('df', 'df_where')
            cp.query('name=="x1"')

        df_ret = cp.run()

        self.assertTrue(df.query('name=="x1"').equals(df_ret))


class DataFrameProxyOperatorTest(unittest.TestCase):

    def test_equalOperator(self):
        cp = CallerProxy()

        with cp.with_cmd() as cmd:
            cmd.create_statement('df', 'cond')
            cp['name'] == 1

        expected = "cond = df['name']==1"
        act = cp.to_code()

        self.assertEqual(expected, act)

    def test_andOperator(self):
        pass


class PandasProxyProxyTest(unittest.TestCase):

    def test_read_excel_code(self):
        import pandas as pd
        px = CallerProxy()
        px.set_obj_to_var('pd', pd)

        with px.with_cmd() as cmd:
            cmd.create_statement('pd')
            px.read_excel('test.xlsx', sheet_name=0)

        expected = "pd.read_excel('test.xlsx',sheet_name = 0)"
        act = px.to_code()

        self.assertEqual(expected, act)

    def test_read_excel_df(self):
        import pandas as pd

        path = 'data/test.xlsx'

        px = CallerProxy()
        px.set_obj_to_var('pd', pd)

        with px.with_cmd() as cmd:
            cmd.create_statement('pd', 'df')
            px.read_excel(path, sheet_name=0)

        df_ret = px.run()
        df_exp = pd.read_excel(path)

        self.assertTrue(df_exp.equals(df_ret))
